import json
from copy import deepcopy  
from langlearn import Logger
from typing import Union, Dict
from langchain.llms.base import LLM
from langlearn.llms import Xinference
def get_llm(name:str) -> Union[LLM|None]:
    # load models.json in current dir and dump to dict
    # {'max_tokens': 256, 'temperature': 0.01, 'top_p': 0.7, 'stream': False, 'stop': ['<|endoftext|>', '<|im_start|>', '<|im_end|>', '<|im_sep|>'], 'stop_token_ids': [2, 6, 7, 8]}
    """
    curl -X 'POST' \
        'http://127.0.0.1:9997/v1/chat/completions' \
        -H 'accept: application/json' \
        -H 'Content-Type: application/json' \
        -d '{
            "model": "yi-chat-34b",
            "messages": [
                {
                    "role": "system",
                    "content": "You are a helpful assistant."
                },
                {
                    "role": "user",
                    "content": "What is the largest animal?"
                }
            ]
        }'

    """
    try:
        with open('models.json', 'r') as file:
            models_dict = json.load(file)
            if not isinstance(models_dict, dict):
                Logger.debug("models json format error")
                return None

            if name not in models_dict:
                Logger.debug(f"{name} not in models.json")
                return None  
            data_copy = deepcopy(models_dict)              
            model_config: Dict = data_copy[name]
            if model_config['provider'] == 'xinference':
                if len(model_config['config']['server_url']) == 0 or len(model_config['config']['model_uid'])==0:
                    Logger.debug("server_url or model_uid not set")
                    return None
                model_kwargs = model_config['config'].get('model_config', {})
                return Xinference(server_url=model_config['config']['server_url'], model_uid=model_config['config']['model_uid'], **model_kwargs)
            else:
                Logger.debug(f"provider {model_config['provider']} not supported")
                return None
    except Exception as e:
        Logger.error(f"load models.json error: {e}")
        return None